import gym
import sys

sys.path.append('./highway-env/scripts/')

sys.path.append('/home/zy/Desktop/RF_decision/RF_decision_demo/highway-env')
import highway_env

sys.path.append('/home/zy/Desktop/RF_decision/RF_decision_demo/rl-agents')
from rl_agents.agents.common.factory import agent_factory

from tqdm import tnrange
from utils import record_videos, show_videos, capture_intermediate_frames


# Make environment
env = gym.make("roundabout-v0")
env = record_videos(env)
obs, done = env.reset(), False
capture_intermediate_frames(env)


# Make agent
agent_config = {
    "__class__": "<class 'rl_agents.agents.tree_search.deterministic.DeterministicPlannerAgent'>",
    "env_preprocessors": [{"method":"simplify"}],
    "budget": 50,
    "gamma": 0.7,
}
agent = agent_factory(env, agent_config)

# Run episode
for step in tnrange(50, desc="Running..."):
    action = agent.act(obs)
    obs, reward, done, info = env.step(action)
    
env.close()
show_videos()